"""
Vision Transformer模型定义
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import timm
from pathlib import Path


class VisionTransformerOOD(nn.Module):
    """用于OOD检测的Vision Transformer"""

    def __init__(self, num_classes=50, model_path=None, freeze_layers=0):
        super().__init__()

        # 创建ViT-L模型
        print("Loading Vision Transformer Large (ViT-L/16)...")
        self.backbone = timm.create_model(
            "vit_large_patch16_224",
            pretrained=False,
            num_classes=0,  # 移除原始分类头
            img_size=224,  # 使用224输入尺寸
        )

        # 加载预训练权重
        if model_path and Path(model_path).exists():
            print(f"Loading weights from {model_path}")
            state_dict = torch.load(model_path, map_location="cpu", weights_only=False)

            # 处理可能的格式问题
            if isinstance(state_dict, dict):
                if "model" in state_dict:
                    state_dict = state_dict["model"]
                elif "state_dict" in state_dict:
                    state_dict = state_dict["state_dict"]

            # 移除分类头权重
            state_dict = {
                k: v
                for k, v in state_dict.items()
                if not k.startswith("head.") and not k.startswith("fc.")
            }

            # 检查是否需要处理位置编码
            if (
                "pos_embed" in state_dict
                and state_dict["pos_embed"].shape != self.backbone.pos_embed.shape
            ):
                print(
                    "Position embedding shape mismatch, keeping model's initialization"
                )
                state_dict.pop("pos_embed")

            # 加载权重
            missing, unexpected = self.backbone.load_state_dict(
                state_dict, strict=False
            )
            print(
                f"Loaded pretrained weights (missing: {len(missing)}, unexpected: {len(unexpected)})"
            )

            # 打印一些缺失的键以便调试
            if missing:
                print(f"Missing keys sample: {missing[:5]}")
        else:
            print(f"ERROR: Pretrained weights not found at {model_path}")
            print("Training ViT from scratch is very difficult and slow!")
            raise FileNotFoundError(
                f"Pretrained weights required but not found at {model_path}"
            )

        # 冻结层（如果需要）
        if freeze_layers > 0:
            self._freeze_layers(freeze_layers)

        # 获取特征维度 - ViT-L是1024
        self.feat_dim = self.backbone.num_features  # 1024 for ViT-L
        print(f"Feature dimension: {self.feat_dim}")

        # 新的分类头
        self.num_classes = num_classes
        self.classifier = nn.Linear(self.feat_dim, num_classes)

        # 初始化分类头 - 使用较大的std
        nn.init.trunc_normal_(self.classifier.weight, std=0.02)
        nn.init.zeros_(self.classifier.bias)

        # 统计参数
        total_params = sum(p.numel() for p in self.parameters())
        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
        print(f"Total parameters: {total_params:,}")
        print(
            f"Trainable parameters: {trainable_params:,} ({100 * trainable_params / total_params:.1f}%)"
        )

        # 打印模型信息
        num_blocks = (
            len(self.backbone.blocks) if hasattr(self.backbone, "blocks") else 24
        )
        print(f"Number of transformer blocks: {num_blocks}")

        self.feature_enhancer = nn.Sequential(
            nn.Linear(self.feat_dim, self.feat_dim),
            nn.ReLU(),
            nn.LayerNorm(self.feat_dim),
            nn.Dropout(0.1)
        )

    def _freeze_layers(self, freeze_layers):
        """冻结前N层"""
        # 冻结patch embedding
        for param in self.backbone.patch_embed.parameters():
            param.requires_grad = False

        # 冻结位置编码
        if hasattr(self.backbone, "pos_embed"):
            self.backbone.pos_embed.requires_grad = False

        # 冻结cls token
        if hasattr(self.backbone, "cls_token"):
            self.backbone.cls_token.requires_grad = False

        # 冻结前N个transformer blocks
        if hasattr(self.backbone, "blocks"):
            for i, block in enumerate(self.backbone.blocks):
                if i < freeze_layers:
                    for param in block.parameters():
                        param.requires_grad = False
                else:
                    for param in block.parameters():
                        param.requires_grad = True

        # 确保norm和head是可训练的
        if hasattr(self.backbone, "norm"):
            for param in self.backbone.norm.parameters():
                param.requires_grad = True

        print(f"Froze {freeze_layers} layers.")

    def forward(self, x, return_features=False):
        """前向传播"""
        # 提取特征
        features = self.backbone(x)

        # 特征增强
        enhanced_features = self.feature_enhancer(features)

        # 特征归一化（可选）
        if return_features:
            features_norm = F.normalize(enhanced_features, p=2, dim=1)
            return features_norm
        # 分类
        logits = self.classifier(enhanced_features)
        return logits
        

    def increment_classes(self, num_new_classes):
        """增加新类别"""
        device = next(self.parameters()).device

        in_features = self.classifier.in_features
        out_features = self.classifier.out_features
        weight = self.classifier.weight.data
        bias = self.classifier.bias.data

        # 创建新的分类头
        self.num_classes += num_new_classes
        self.classifier = nn.Linear(in_features, self.num_classes)
        self.classifier.to(device)

        # 复制旧类的权重
        self.classifier.weight.data[:out_features] = weight
        self.classifier.bias.data[:out_features] = bias

        # 初始化新类的权重 - 使用更好的策略
        # 1. 计算旧类权重的统计信息
        old_weight_norm = weight.norm(dim=1).mean()
        old_weight_std = weight.std().item()

        # 2. 使用相似的分布初始化新类权重
        new_weight = torch.zeros((num_new_classes, in_features), device=device)
        nn.init.normal_(new_weight, mean=0, std=old_weight_std)
        self.classifier.weight.data[out_features:] = new_weight

        # 3. 调整新类权重的范数
        for i in range(out_features, self.num_classes):
            current_norm = self.classifier.weight.data[i].norm()
            if current_norm > 0:
                self.classifier.weight.data[i] = self.classifier.weight.data[i] * (
                    old_weight_norm / current_norm
                )

        # 4. 初始化新类偏置 - 使用旧类偏置的平均值作为起点
        old_bias_mean = bias.mean().item()
        nn.init.constant_(self.classifier.bias.data[out_features:], old_bias_mean)

        print(f"Extended classifier from {out_features} to {self.num_classes} classes")
        print(
            f"Old class weight norm: {old_weight_norm:.4f}, std: {old_weight_std:.4f}"
        )
        print(
            f"New class weight norm: {self.classifier.weight.data[out_features:].norm(dim=1).mean():.4f}"
        )
        print(f"Old bias mean: {old_bias_mean:.4f}")
        print(
            f"New bias initialized to: {self.classifier.bias.data[out_features:].mean():.4f}"
        )


class IncrementalLearner:
    """增量学习器"""

    def __init__(self, model, old_model=None):
        self.model = model
        self.old_model = old_model

    def compute_distillation_loss(self, outputs, old_outputs, temperature):
        """计算知识蒸馏损失"""
        # 只对旧类进行蒸馏
        num_old_classes = old_outputs.size(1)

        # 软目标
        soft_targets = F.softmax(old_outputs / temperature, dim=1)
        soft_outputs = F.log_softmax(outputs[:, :num_old_classes] / temperature, dim=1)

        # KL散度损失
        distillation_loss = F.kl_div(
            soft_outputs, soft_targets, reduction="batchmean"
        ) * (temperature**2)

        return distillation_loss

    def balanced_fine_tuning(self, outputs, targets, num_old_classes):
        """平衡微调：对新旧类使用不同的损失权重"""
        # 分离新旧类
        old_classes_mask = targets < num_old_classes
        new_classes_mask = ~old_classes_mask

        loss = 0

        # 旧类损失
        if old_classes_mask.any():
            old_outputs = outputs[old_classes_mask]
            old_targets = targets[old_classes_mask]
            old_loss = F.cross_entropy(old_outputs, old_targets)
            loss += old_loss * 0.5  # 降低旧类权重

        # 新类损失
        if new_classes_mask.any():
            new_outputs = outputs[new_classes_mask]
            new_targets = targets[new_classes_mask]
            new_loss = F.cross_entropy(new_outputs, new_targets)
            loss += new_loss * 1.5  # 增加新类权重

        return loss
